

#include "trt.hpp"


// 预处理输入图像
std::vector<float> preprocessImage(const cv::Mat& image, int inputHeight, int inputWidth)
{
	// 检查图像是否为空
	if (image.empty())
	{
		throw std::runtime_error("输入图像为空");
	}

	// 调整图像大小并转换为浮点类型
	cv::Mat img;
	cv::resize(image, img, cv::Size(inputWidth, inputHeight)); // 调整图像大小
	img.convertTo(img, CV_32FC3, 1.0f / 255.0f); // 再进行类型转换和归一化

	// 归一化 (使用均值和标准差)
	cv::Scalar mean(0.485f, 0.456f, 0.456f);
	cv::Scalar std(0.229f, 0.224f, 0.225f);
	img = (img - mean) / std;

	// 将数据从 HWC 转换为 NCHW
	std::vector<float> inputData(inputHeight * inputWidth * 3);
	for (int c = 0; c < 3; ++c)
	{
	    for (int h = 0; h < inputHeight; ++h)
	    {
		for (int w = 0; w < inputWidth; ++w)
		{
		    // 正确的方式是使用 at 成员函数来访问像素值，并使用 cv::Vec3f 作为模板参数
		    inputData[c * inputHeight * inputWidth + h * inputWidth + w] = img.at<cv::Vec3f>(h, w)[c];
		}
	    }
	}

	return inputData;
}

void* InferenceThread(void* pArgs) {
    // 从传入的用户数据中获取线程参数
    InferenceThreadArgs* args = static_cast<InferenceThreadArgs*>(pArgs);

    // 检查参数有效性
    if (args == nullptr || args->trtModel == nullptr || args->sharedResWrapper == nullptr) {
        std::cerr << "Invalid thread arguments provided." << std::endl;
        return nullptr;
    }
    // 从参数中获取 TensorRT 模型和共享资源
    TensorRTModel* trtModel = args->trtModel;
    SharedResourcesWrapper* resources = args->sharedResWrapper;

    while (!g_bExit) {
        // 使用 std::unique_lock 锁定 FrameSharedResources 的互斥锁，确保线程安全。
        std::unique_lock<std::mutex> lock(resources->frameRes->mutex);

        // 等待条件变量，直到有新的帧可用或 g_bExit 被设置为 true。
        resources->frameRes->cond.wait(lock, [&resources] { return resources->frameRes->newFrameAvailable || g_bExit; });

        if (g_bExit) {
            std::cout << "Exit signal received in InferenceThread." << std::endl;
            break;
        }

        // 检查 FrameSharedResources 中是否有待处理的帧。
        if (!resources->frameRes->frames.empty()) {
            // 获取共享资源中最后一个图像帧的引用。
            ImageFrame& lastFrame = resources->frameRes->frames.back();
            cv::Mat originalImage = *lastFrame.mat; // 获取图像数据
            int frameNum = lastFrame.frameNum;       // 获取帧编号
            std::time_t timestamp = lastFrame.timestamp; // 获取时间戳

            // 解锁互斥锁，因为图像处理和推理可能较耗时，避免长时间占用锁。
            lock.unlock();

            // 预处理图像，准备输入到推理模型。
            auto inputData = preprocessImage(originalImage, 224, 224);

            // 执行推理操作。
            auto outputData = trtModel->infer(inputData);

            // 根据推理结果计算预测类别。
            int predictedClass = (outputData[0] > outputData[1]) ? 1 : 0; // 假设 outputData[0] 和 outputData[1] 是模型的输出

            // 创建新的推理结果实例
            InferenceResult newInferenceResult;
            newInferenceResult.mat = std::make_shared<cv::Mat>(originalImage.clone()); // 存储图像副本
            newInferenceResult.frameNum = frameNum; // 设置帧编号
            newInferenceResult.timestamp = timestamp; // 设置时间戳
            newInferenceResult.predictedClass = predictedClass; // 设置预测类别

            // 将推理结果添加到推理结果共享资源。
            addInferenceResultToSharedResources(newInferenceResult, *resources->inferRes);

            // 锁定互斥锁，准备移除已处理的帧。
            lock.lock();
            // 从共享资源中移除已处理的帧，确保不会重复处理。
            resources->frameRes->frames.pop_back();
            resources->frameRes->newFrameAvailable = false; // 重置新帧可用标志
        } else {
            std::cout << "No new frames to process." << std::endl;
        }
    }

    // 线程函数结束，返回 nullptr。
    std::cout << "InferenceThread has finished processing." << std::endl;
    return nullptr;
}

// 确保addInferenceResultToSharedResources函数与InferenceSharedResources兼容
void addInferenceResultToSharedResources(InferenceResult& result, InferenceSharedResources& sharedRes) {
    // 锁定互斥锁以安全地修改共享资源
    std::lock_guard<std::mutex> lock(sharedRes.mutex);

    // 将新的推理结果添加到共享资源的向量中
    sharedRes.infer_frame.push_back(result);

    // 标记有新的推理结果可用
    sharedRes.newInferenceResultAvailable = true;

    // 通知等待的线程有新的推理结果
    sharedRes.cond.notify_one();
}

// 用于显示通知的函数
void showNotification(const std::string& title, const std::string& message) {
    std::string command =
	 "notify-send -u critical -i 'notification-error' '" + title + "' '" + message + "'";
    system(command.c_str());
}

// 函数定义
int readThresholdFromFile(const std::string& filename) {
    std::ifstream file(filename);
    if (!file) {
        throw std::runtime_error("无法打开文件: " + filename);
    }

    int threshold;
    file >> threshold;
    if (file.fail()) {
        throw std::runtime_error("读取文件时发生错误: " + filename);
    }
	
    file.close();
    return threshold;
}



void drawCircle(int radius) {
    const char* circle = "*";
    const int circleLength = 3; // 圆形字符的长度
    int diameter = 2 * radius + 1; // 圆形的直径
    int circleRows = (diameter + circleLength - 1) / circleLength;

    for (int i = 0; i < circleRows; ++i) {
        for (int j = 0; j < diameter; ++j) {
            double x = j - radius;
            double y = i - radius;
            double distance = std::sqrt(x * x + y * y);
            if (distance < radius + 0.5) {
                std::cout << circle[1];
            } else {
                std::cout << " ";
            }
        }
        std::cout << std::endl;
    }
}

void* InferDisplayThread(void* pUser) {
    // 从传入的用户数据中获取推理结果共享资源。
    InferenceSharedResources* inferRes = static_cast<InferenceSharedResources*>(pUser);

    if (!inferRes) {
        return nullptr;
    }

    std::unique_lock<std::mutex> lock(inferRes->mutex); // 锁定互斥锁

    int threshold = readThresholdFromFile("/home/jetson/Documents/trt_3d/thres.txt"); // 假设从某处获取的阈值
    int countClassOne = 0; // 累计predictedClass为1的数量
    bool shouldDrawStars = false;

    while (!g_bExit) {
        inferRes->cond.wait(lock, [&inferRes] { return inferRes->newInferenceResultAvailable || g_bExit; });

        if (g_bExit) {
            break;
        }

         while (!inferRes->infer_frame.empty()) {
            InferenceResult result = inferRes->infer_frame.front();
            inferRes->infer_frame.erase(inferRes->infer_frame.begin());

            if (result.predictedClass == 1) {
                countClassOne++; // 累计数量
                std::cout << "Inference complete. Predicted class: " << result.predictedClass << std::endl;

                if (countClassOne > threshold) {
                    shouldDrawStars = true;
                    break; // 跳出循环，准备画星号
                }
            } 
        }

        if (shouldDrawStars) {
            // 显示四行星号
            for (int i = 0; i < 4; ++i) {
                std::cout << "**************************************************" << std::endl;
            }
            std::cout << std::endl; // 添加一个空行
        }

        inferRes->newInferenceResultAvailable = false;
    }

    return nullptr;
}

std::vector<cv::Mat> loadImages(const std::string& directoryPath) {
    std::vector<cv::Mat> images;
    std::string pattern = directoryPath + "/*.jpg"; // 假设图像文件是jpg格式，你可以修改为其他格式如png等
    std::vector<std::string> filenames;

    // 使用cv::glob获取所有匹配的文件名
    cv::glob(pattern, filenames);

    // 遍历文件名列表，读取每个图像
    for (const auto& filename : filenames) {
        cv::Mat image = cv::imread(filename, cv::IMREAD_COLOR); // 读取彩色图像
        if (!image.empty()) {
            images.push_back(image); // 将图像添加到向量中
        } else {
            std::cerr << "Could not read image: " << filename << std::endl;
        }
    }

    return images;
}
